# -----------------------------------------------------------
# 2020/10/19
# 读取CloudSat 2B-GEOPROF数据,画出卫星轨迹和雷达反射率因子的剖面图.
# 参考:
# [1] https://hdfeos.org/zoo/index_openCDPC_Examples.php#CloudSAT
# [2] https://moonbooks.org/Articles/How-to-read-CloudSat-2B-GEOPROF-GRANULE-HDF4-file-using-python-and-pyhdf-/
# -----------------------------------------------------------
from read_CloudSat import reader

import numpy as np
import numpy.ma as ma

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

def set_map(ax):
    '''设置地图.'''

    proj = ccrs.PlateCarree()
    ax.coastlines(lw=0.5)

    xticks = np.arange(-180, 180 + 60, 60)
    yticks = np.arange(-90, 90 + 30, 30)
    ax.set_xticks(xticks, crs=proj)
    ax.set_yticks(yticks, crs=proj)
    ax.xaxis.set_major_formatter(LongitudeFormatter())
    ax.yaxis.set_major_formatter(LatitudeFormatter())
    ax.tick_params('both', labelsize='x-small')

    ax.set_global()


def draw_track(ax, lon1D, lat1D):
    '''根据经纬度画出轨迹,并标识出起点.'''

    ax.plot(lon1D, lat1D, lw=2, color='b', transform=ccrs.Geodetic())
    ax.plot(lon1D[0], lat1D[0], 'ro', ms=3, transform=ccrs.PlateCarree())
    ax.text(
        lon1D[0] + 5, lat1D[0], 'start', color='r', fontsize='x-small',
        transform=ccrs.PlateCarree()
    )
def draw_cross_section(ax, lon, hgt, data):
    '''画出data的经度-高度剖面.'''
    im = ax.pcolormesh(lon, hgt, data.T, cmap='jet', shading='nearest')
    ax.set_ylim(0, 20)
    ax.set_xlim(lon.min(), lon.max())  # 设置经度范围
    ax.tick_params(labelsize='x-small')
    ax.set_xlabel('Longitude', fontsize='x-small')  # 改横坐标标签为经度
    ax.set_ylabel('Height [km]', fontsize='x-small')

    divider = make_axes_locatable(ax)
    cax = divider.append_axes('bottom', size='10%', pad=0.42)
    cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
    cbar.ax.tick_params(labelsize='x-small')
    cbar.set_label('dBZe', fontsize='x-small')

def draw_elevation(ax, lon, elv):
    '''根据经度画出地形抬升.'''
    ax.fill_between(lon, elv, color='gray')

if __name__ == '__main__':
    # 读取数据（原有代码不变）
    fname = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_GEOPROF/cloudsat_2020_1_4/2020112031512_74484_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
    f = reader(fname)
    lon, lat, elv = f.read_geo()
    data = f.read_sds('Radar_Reflectivity')
    height = f.read_sds('Height')
    time = f.read_time(datetime=False)
    start_time, end_time = f.read_time(datetime=True)[[0, -1]]
    f.close()

    hgt = height.mean(axis=0).data / 1000  # 高度转km
    elv /= 1000  # 地形高度转km

    fig = plt.figure(dpi=200, figsize=(6, 6))
    ax1 = fig.add_axes([0.25, 0.45, 0.5, 0.5], projection=ccrs.PlateCarree())
    ax2 = fig.add_axes([0.25, 0.2, 0.5, 0.3])

    set_map(ax1)
    draw_track(ax1, lon, lat)
    ax1.set_title(f'From {start_time.strftime("%Y-%m-%d %H:%M")} to {end_time.strftime("%Y-%m-%d %H:%M")}', fontsize='small')

    # 修改调用，传入lon而非time
    draw_cross_section(ax2, lon, hgt, data)
    draw_elevation(ax2, lon, elv)
    ax2.set_title('Radar Reflectivity Factor', fontsize='small')

    fig.savefig('test2', bbox_inches='tight')
    plt.show()
    plt.close(fig)
#
# def draw_cross_section(ax, time, hgt, data):
#     '''画出data的时间-高度剖面.'''
#
#     im = ax.pcolormesh(time, hgt, data.T, cmap='jet', shading='nearest')
#
#     ax.set_ylim(0, 20)
#     ax.set_xlim(0, time.max())
#     ax.tick_params(labelsize='x-small')
#     ax.set_xlabel('Seconds since the start of the granule [s]', fontsize='x-small')
#     ax.set_ylabel('Height [km]', fontsize='x-small')
#
#     divider = make_axes_locatable(ax)
#     cax = divider.append_axes('bottom', size='10%', pad=0.42)
#     cbar = plt.colorbar(im, cax=cax, extend='both', orientation='horizontal')
#     cbar.ax.tick_params(labelsize='x-small')
#     cbar.set_label('dBZe', fontsize='x-small')
#
#
# def draw_elevation(ax, time, elv):
#     '''画出地形抬升.'''
#
#     ax.fill_between(time, elv, color='gray')
#
#
# if __name__ == '__main__':
#     # 读取数据
#     # fname = '2019001023402_67527_CS_2B-GEOPROF_GRANULE_P1_R05_E08_F03.hdf'
#     fname = '/mnt/datastore/liudddata/cloudsat_data/cloudsat_GEOPROF/cloudsat_2020_1_4/2020112031512_74484_CS_2B-GEOPROF_GRANULE_P1_R05_E09_F00.hdf'
#     f = reader(fname)
#     lon, lat, elv = f.read_geo()
#     data = f.read_sds('Radar_Reflectivity')
#     height = f.read_sds('Height')
#     time = f.read_time(datetime=False)
#     start_time, end_time = f.read_time(datetime=True)[[0, -1]]
#     f.close()
#     fr = open("/home/liudd/code/cloudsat_process/write_word.txt", "w", encoding="UTF-8")
#    # fr.write(lon)
#     fr.close
#     # 每个profile的高度有微小的差别,这里用平均值代替所有profile的高度.
#     hgt = height.mean(axis=0).data
#     # 高度单位全部转为km.
#     elv /= 1000
#     hgt /= 1000
#
#     fig = plt.figure(dpi=200, figsize=(6, 6))
#     ax1 = fig.add_axes([0.25, 0.45, 0.5, 0.5], projection=ccrs.PlateCarree())
#     ax2 = fig.add_axes([0.25, 0.2, 0.5, 0.3])
#
#     set_map(ax1)
#     draw_track(ax1, lon, lat)
#
#     # 把granule的起止时间标上.
#     start_str = start_time.strftime('%Y-%m-%d %H:%M')
#     end_str = end_time.strftime('%Y-%m-%d %H:%M')
#     ax1.set_title(f'From {start_str} to {end_str}', fontsize='small')
#
#     draw_cross_section(ax2, time, hgt, data)
#     draw_elevation(ax2, time, elv)
#     ax2.set_title('Radar Reflectivity Factor', fontsize='small')
#
#     fig.savefig('test2', bbox_inches='tight')
#     print(type(lon))
#     plt.show()
#     plt.close(fig)
